import os
import json
from typing import Dict, List, Any, Optional, Tuple
from pydantic import BaseModel, Field
from loguru import logger
from dotenv import load_dotenv
import traceback
from datetime import datetime, date

from .db import get_db_operations
from .model import get_model_client
from .utils import validate_sql, format_results, CustomJSONEncoder

class QueryRequest(BaseModel):
    """自然语言查询请求"""
    query: str = Field(..., description="自然语言查询")
    format: str = Field("text", description="结果格式: text, json, html, csv")
    include_explanation: bool = Field(True, description="是否包含SQL解释")
    include_analysis: bool = Field(True, description="是否包含结果分析")
    include_visualization: bool = Field(True, description="是否包含可视化建议")

class SQLExecuteRequest(BaseModel):
    """SQL执行请求"""
    sql: str = Field(..., description="SQL查询语句")
    format: str = Field("text", description="结果格式: text, json, html, csv")
    validate: bool = Field(True, description="是否验证SQL")

class CoreProcessor:
    """核心处理器，处理客户端请求并集成各个组件"""
    
    def __init__(self):
        """初始化处理器"""
        self.db = get_db_operations()
        self.model = get_model_client()
        logger.info("核心处理器已初始化")
        
        # 缓存数据库元数据
        self._metadata = None
        self._example_data = {}
        self._db_connected = False
    
    def _ensure_db_connection(self):
        """确保数据库连接已建立"""
        if not self._db_connected:
            try:
                # 重新加载环境变量
                load_dotenv(override=True)
                
                # 重新连接数据库
                self.db.db.connect()
                self._db_connected = True
                logger.info(f"已连接到数据库: {os.getenv('DB_HOST', 'localhost')}")
            except Exception as e:
                logger.error(f"数据库连接失败: {str(e)}")
                raise
    
    def _get_database_metadata(self, refresh: bool = False) -> Dict[str, Any]:
        """获取数据库元数据，可以选择刷新缓存"""
        # 确保数据库连接
        self._ensure_db_connection()
        
        if self._metadata is None or refresh:
            logger.info("正在获取数据库元数据...")
            try:
                # 获取元数据
                self._metadata = self.db.get_db_metadata()
                
                # 元数据检查
                if not self._metadata:
                    logger.warning("获取到的数据库元数据为空")
                else:
                    logger.success(f"已获取元数据: {len(self._metadata)} 张表")
                
            except Exception as e:
                logger.error(f"获取数据库元数据失败: {str(e)}")
                self._metadata = {}
                # 重新抛出异常以便上层函数捕获
                raise ValueError(f"无法获取数据库元数据: {str(e)}")
                
        return self._metadata
    
    def _get_sample_data(self, table_name: str) -> List[Dict[str, Any]]:
        """获取表的样本数据"""
        if table_name not in self._example_data:
            try:
                df, error = self.db.get_sample_data(table_name, limit=3)
                if error or df is None:
                    logger.warning(f"获取表 {table_name} 的样本数据失败: {error}")
                    self._example_data[table_name] = []
                else:
                    # 安全转换DataFrame为字典列表
                    try:
                        # 处理日期时间类型
                        for col in df.select_dtypes(include=['datetime64', 'datetime64[ns]']).columns:
                            df[col] = df[col].dt.strftime('%Y-%m-%d %H:%M:%S')
                        
                        records = df.to_dict(orient="records")
                        # 额外检查每个值的类型
                        for record in records:
                            for key, value in list(record.items()):
                                if isinstance(value, (datetime, date)):
                                    record[key] = value.isoformat()
                                elif not isinstance(value, (str, int, float, bool, type(None))):
                                    record[key] = str(value)
                        
                        self._example_data[table_name] = records
                    except Exception as e:
                        logger.error(f"转换表 {table_name} 的样本数据时出错: {str(e)}")
                        self._example_data[table_name] = []
            except Exception as e:
                logger.error(f"获取表 {table_name} 的样本数据时出错: {str(e)}")
                self._example_data[table_name] = []
        return self._example_data[table_name]
    
    def process_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """处理客户端请求"""
        action = request.get("action", "").lower()
        
        if action == "query":
            return self.handle_nl_query(request)
        elif action == "sql":
            return self.handle_sql_query(request)
        elif action == "info":
            return self.handle_db_info(request)
        else:
            logger.warning(f"未知的请求操作: {action}")
            return {"error": f"未支持的操作: {action}"}
    
    def handle_nl_query(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """处理自然语言查询"""
        # 验证请求参数
        try:
            req = QueryRequest(**request)
        except Exception as e:
            return {"error": f"无效的请求格式: {str(e)}"}
        
        # 获取数据库元数据
        try:
            metadata = self._get_database_metadata()
            if not metadata:
                return {"error": "无法获取数据库元数据"}
        except Exception as e:
            return {"error": f"获取数据库元数据失败: {str(e)}"}
        
        # 为关键表获取样本数据
        example_data = {}
        for table_name in metadata.keys():
            try:
                example_data[table_name] = self._get_sample_data(table_name)
            except Exception as e:
                logger.error(f"获取表 {table_name} 的样本数据失败: {str(e)}")
                example_data[table_name] = []
        
        logger.info(f"处理自然语言查询: {req.query}")
        
        try:
            # 使用模型生成SQL
            sql_result = self.model.generate_sql(
                req.query, metadata, example_data
            )
            
            # 检查是否有错误
            if "error" in sql_result and sql_result["error"]:
                logger.error(f"生成SQL失败: {sql_result['error']}")
                return {
                    "error": f"生成SQL失败: {sql_result['error']}",
                    "query": req.query
                }
            
            generated_sql = sql_result.get("generated_sql", "")
            if not generated_sql:
                logger.error("生成的SQL为空")
                return {
                    "error": "生成的SQL为空",
                    "query": req.query
                }
                
            logger.info(f"生成的SQL: {generated_sql}")
            
            # 验证生成的SQL
            validation = validate_sql(generated_sql)
            if not validation["is_valid"]:
                return {
                    "error": "生成的SQL无效",
                    "validation": validation,
                    "generated_sql": generated_sql
                }
            
            if validation["warnings"]:
                logger.warning(f"SQL警告: {', '.join(validation['warnings'])}")
            
            # 执行SQL查询
            try:
                query_results, error = self.db.execute_query(generated_sql)
                
                if error:
                    return {
                        "error": f"执行查询失败: {error}",
                        "generated_sql": generated_sql,
                        "validation": validation
                    }
                
                # 确保查询结果是可序列化的
                safe_results = []
                for row in query_results:
                    safe_row = {}
                    for key, value in row.items():
                        if isinstance(value, (datetime, date)):
                            safe_row[key] = value.isoformat()
                        elif value is not None and not isinstance(value, (str, int, float, bool)):
                            safe_row[key] = str(value)
                        else:
                            safe_row[key] = value
                    safe_results.append(safe_row)
                
                # 格式化结果
                formatted_results = format_results(safe_results, req.format)
                
                response = {
                    "query": req.query,
                    "generated_sql": generated_sql,
                    "results": safe_results,
                    "formatted_results": formatted_results,
                    "validation": validation
                }
                
                # 添加额外信息（如果请求）
                if req.include_explanation:
                    try:
                        explanation = self.model.explain_sql(generated_sql)
                        response["explanation"] = explanation
                    except Exception as e:
                        logger.error(f"解释SQL时出错: {str(e)}")
                        response["explanation"] = f"无法解释SQL: {str(e)}"
                
                if req.include_analysis and safe_results:
                    try:
                        analysis = self.model.analyze_query_results(safe_results, req.query)
                        response["analysis"] = analysis
                    except Exception as e:
                        logger.error(f"分析查询结果时出错: {str(e)}")
                        response["analysis"] = f"无法分析结果: {str(e)}"
                
                if req.include_visualization and safe_results:
                    try:
                        viz_suggestion = self.model.suggest_visualization(safe_results, req.query)
                        response["visualization"] = viz_suggestion
                    except Exception as e:
                        logger.error(f"获取可视化建议时出错: {str(e)}")
                        response["visualization"] = f"无法获取可视化建议: {str(e)}"
                
                return response
                
            except Exception as e:
                logger.error(f"执行SQL查询时出错: {str(e)}")
                return {
                    "error": f"执行查询失败: {str(e)}",
                    "generated_sql": generated_sql,
                    "validation": validation
                }
        
        except Exception as e:
            logger.error(f"处理自然语言查询时出错: {str(e)}")
            logger.error(f"堆栈跟踪: {traceback.format_exc()}")
            return {"error": f"处理查询失败: {str(e)}"}
    
    def handle_sql_query(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """处理SQL查询请求"""
        try:
            req = SQLExecuteRequest(**request)
        except Exception as e:
            return {"error": f"无效的请求格式: {str(e)}"}
        
        logger.info(f"执行SQL查询: {req.sql}")
        
        # 验证SQL（如果请求）
        validation = None
        if req.validate:
            validation = validate_sql(req.sql)
            if not validation["is_valid"]:
                return {
                    "error": "SQL验证失败",
                    "validation": validation
                }
            
            if validation["warnings"]:
                logger.warning(f"SQL警告: {', '.join(validation['warnings'])}")
        
        # 执行查询
        try:
            if validation and validation["query_type"] in ["select"]:
                query_results, error = self.db.execute_query(req.sql)
            else:
                success, error = self.db.execute_write_query(req.sql)
                if success:
                    query_results = [{"affected_rows": "操作成功完成"}]
                else:
                    query_results = []
            
            if error:
                return {
                    "error": f"执行SQL失败: {error}",
                    "validation": validation
                }
            
            # 确保查询结果是可序列化的
            safe_results = []
            for row in query_results:
                safe_row = {}
                for key, value in row.items():
                    if isinstance(value, (datetime, date)):
                        safe_row[key] = value.isoformat()
                    elif value is not None and not isinstance(value, (str, int, float, bool)):
                        safe_row[key] = str(value)
                    else:
                        safe_row[key] = value
                safe_results.append(safe_row)
            
            # 格式化结果
            formatted_results = format_results(safe_results, req.format)
            
            return {
                "sql": req.sql,
                "results": safe_results,
                "formatted_results": formatted_results,
                "validation": validation
            }
        except Exception as e:
            logger.error(f"执行SQL查询时出错: {str(e)}")
            logger.error(f"堆栈跟踪: {traceback.format_exc()}")
            return {
                "error": f"执行SQL失败: {str(e)}",
                "validation": validation
            }
    
    def handle_db_info(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """处理数据库信息请求"""
        try:
            # 获取数据库元数据
            metadata = self._get_database_metadata(refresh=True)
            
            # 如果元数据为空，则返回错误
            if not metadata:
                return {"error": "无法获取数据库元数据，请检查数据库连接"}
            
            # 获取表的行数
            table_counts = {}
            for table_name in metadata.keys():
                try:
                    count = self.db.count_rows(table_name)
                    table_counts[table_name] = count
                except Exception as e:
                    logger.error(f"获取表 {table_name} 的行数时出错: {str(e)}")
                    table_counts[table_name] = "无法获取"
            
            # 构建响应
            response = {
                "tables": list(metadata.keys()),
                "table_counts": table_counts,
                "metadata": metadata
            }
            
            # 将响应序列化为JSON，处理特殊类型
            try:
                # 使用自定义JSON编码器，确保所有对象可序列化
                return json.loads(json.dumps(response, cls=CustomJSONEncoder))
            except Exception as e:
                logger.error(f"序列化数据库信息时出错: {str(e)}")
                return {"error": f"无法序列化数据库信息: {str(e)}"}
            
        except Exception as e:
            logger.error(f"处理数据库信息请求时出错: {str(e)}")
            logger.error(f"堆栈跟踪: {traceback.format_exc()}")
            return {"error": f"获取数据库信息失败: {str(e)}"}
    
    def close(self):
        """关闭处理器及相关资源"""
        try:
            self.db.close()
            logger.info("处理器资源已释放")
        except Exception as e:
            logger.error(f"关闭处理器资源时出错: {str(e)}")


# 单例模式
_processor = None

def get_processor() -> CoreProcessor:
    """获取核心处理器的单例实例"""
    global _processor
    if _processor is None:
        _processor = CoreProcessor()
    return _processor 